#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/8 17:26
# @Author  : jeff
# @FileName: security_aspect
import functools

import jwt
from flask import request, jsonify, make_response
from src.model.role_model import Role
from src.model.user_model import User
from src.utils.jwt_utils import JwtUtils
from src.enums.status_enum import Status


def require_token(role=None):
    """
    Verification token
     1. Obtain the token from the request header. If the token does not exist, an error will be returned.
     2. If the token exists, verify whether the token is valid. If it is invalid, return an error.
     3. If the token is valid, check whether there is permission. If there is no permission, return an error.
     4. If you have permission, renew the token and return the token in the response header.
    :param role:
    :return:
    """

    def decorator(f):
        @functools.wraps(f)
        def decorated_function(*args, **kwargs):
            token = request.headers.get('token')
            if not token:
                return make_response(jsonify(Status.MISS_TOKEN.value))
            try:
                payload = JwtUtils.get_payload(token)
            except jwt.exceptions.ExpiredSignatureError:
                return make_response(jsonify(Status.INVALID_TOKEN.value))
            except jwt.exceptions.InvalidSignatureError:
                return make_response(jsonify(Status.VERIFICATION_FAILED.value))
            except jwt.exceptions.DecodeError:
                return make_response(jsonify(Status.VERIFICATION_FAILED.value))
            if not check_admin_permission(payload, role):
                return make_response(jsonify(Status.PERMISSION_FAILED.value))
            new_token = JwtUtils.renew_token(token=token)
            response = f(*args, **kwargs)
            response.headers['token'] = new_token
            return response

        return decorated_function

    return decorator


def check_admin_permission(payload, role):
    users = User.query_by_user(
        username=payload['username'],
        password=payload['password']
    )

    if role is None:
        return True
    for user in users:
        user_role = Role.query_by_id(role_id=user.role_id)
        if user_role:
            if role == user_role[0].role_name:
                return True
        else:
            return False
    return False
